package main

import (
	"context"
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

var (
	cn  = flag.Int("c", 1, "connection number, default 1 connection")
	bd  = flag.String("base-dir", ".", "base directory, default is current directory")
	d   = flag.Duration("d", time.Second, "benchmark duration default time.Second")
	uri = flag.String("uri", "", "benchmark uri")

	validResp = flag.String("resp", "", "validate response content, default not valid")

	wg  = &sync.WaitGroup{}
	cli = &Client{}

	reqTotal   uint64 = 0
	reqSuc     uint64 = 0
	reqErr     uint64 = 0
	reqRespErr uint64 = 0

	errNotRedirect = errors.New("it doesn't follow redirect")
)

// go run bench_client.go -c=100 -base-dir=. -uri=http://test.local -d=10s
// bench_client -c=100 -base-dir=. -uri=http://test.local -d=10s
func main() {
	flag.Parse()
	if *uri == "" {
		flag.PrintDefaults()
		log.Fatal("uri is required")
		return
	}

	logPath := strings.TrimRight(*bd, "/") + "/" + time.Now().Format("2006-01-02_150405") + ".log"

	f, err := os.OpenFile(logPath, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		fmt.Println(err)
		return
	}

	defer func() { _ = f.Close() }()

	log.SetOutput(f)

	ctx, cancel := context.WithTimeout(context.Background(), *d)
	defer cancel()

	for i := 0; i < *cn; i++ {
		wg.Add(1)
		go get(ctx, *uri)
	}

	wg.Wait()

	fmt.Println("send req total:", reqTotal)
	fmt.Println("send req success:", reqSuc)
	fmt.Println("send req err:", reqErr)
	fmt.Println("send req response err:", reqRespErr)
}

func ParseUrl(reqUrl string, params url.Values) (*url.URL, error) {
	reqUrlObj, err := url.Parse(reqUrl)
	if err != nil {
		return reqUrlObj, err
	}
	urlParams := reqUrlObj.Query()
	if params != nil {
		for k, v := range params {
			urlParams[k] = v
		}
		reqUrlObj.RawQuery = urlParams.Encode()
	}

	return reqUrlObj, nil
}

func get(ctx context.Context, uri string) error {
	defer wg.Done()

	for {
		select {
		case <-ctx.Done():
			return nil
		default:
			atomic.AddUint64(&reqTotal, 1)
			body, _, resp, err := cli.Get(uri, url.Values{}, nil)
			if err != nil {
				atomic.AddUint64(&reqErr, 1)
				log.Println(err)
				continue
			}

			if resp.StatusCode != http.StatusOK {
				atomic.AddUint64(&reqErr, 1)
				log.Printf("http status:%d", resp.StatusCode)
				continue
			}

			if *validResp != "" && string(body) != *validResp {
				atomic.AddUint64(&reqRespErr, 1)
				log.Println("response body is invalid")
				continue
			}

			atomic.AddUint64(&reqSuc, 1)
			log.Println(string(body))
		}
	}

}

type Client struct {
	proxyUrl    string
	cli         *http.Client
	notRedirect bool
}

func (c Client) Get(reqUrl string, params url.Values, headers map[string]string) (bodyBytes []byte, req *http.Request, resp *http.Response, err error) {
	reqHeader := http.Header{}
	for k := range headers {
		reqHeader.Set(k, headers[k])
	}
	req, err = http.NewRequest(http.MethodGet, reqUrl, nil)
	if err != nil {
		return nil, req, nil, err
	}
	req.URL, err = ParseUrl(reqUrl, params)
	if err != nil {
		return nil, req, nil, err
	}

	return c.httpClient(req, nil, reqHeader)
}

func (c Client) getClient() *http.Client {
	if c.cli != nil {
		return c.cli
	}
	if c.proxyUrl != "" {
		proxy, _ := url.Parse(c.proxyUrl)
		tr := &http.Transport{
			Proxy:           http.ProxyURL(proxy),
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}

		c.cli = &http.Client{
			Transport: tr,
			Timeout:   time.Second * 5, // 超时时间
		}
	} else {
		c.cli = &http.Client{}
	}
	if c.notRedirect {
		c.cli.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return errNotRedirect
		}
	}

	return c.cli
}

func (c Client) httpClient(httpReq *http.Request, reqBody io.Reader, headers http.Header) (bodyBytes []byte, req *http.Request, resp *http.Response, err error) {
	rc, ok := reqBody.(io.ReadCloser)
	if !ok && reqBody != nil {
		rc = io.NopCloser(reqBody)
	}
	httpReq.Body = rc
	httpReq.Header = headers

	resp, err = c.getClient().Do(httpReq)
	// Not follow redirect err Skip
	if err != nil && !strings.Contains(err.Error(), errNotRedirect.Error()) {
		return nil, httpReq, resp, err
	}
	defer func() { _ = resp.Body.Close() }()

	bodyBytes, err = io.ReadAll(resp.Body)
	if err != nil {
		return nil, httpReq, resp, err
	}

	return bodyBytes, httpReq, resp, err
}
